# Ofri Keidar Inbal Wiesel 302933981 305331878

'''
Class implements Lidstone smoothing for bigram model
'''
class BiLidstone:
     
    '''
    Constructor- creates a new object implementing
    Lidstone smoothing for bigram model
    '''
    def __init__(self, VOCAB_SIZE, uniObsVoc, biObsVoc, lambdaVal, numArticles):
        
        # set members
        self.VOCAB_SIZE = VOCAB_SIZE; # language vocabulary size
        self.uniObsVoc = uniObsVoc; # maps each word (unigram) to its frequency in train set
        self.obsVoc = biObsVoc; # maps each bigram to its frequency in train set
        self.lambdaVal = lambdaVal; # value of lambda smoothing parameter
        self.BEGIN_ARTICLE = "begin-article-event"; # virtual event at article's beginning
        self.numArticles = numArticles; # frequency of virtual "begin-article" unigram event 
    
    '''
    Calculates the discounted probability for given bigram, using Lidstone smoothing
    '''
    def getDiscProb(self, inputBigram):
        
        # calculate numerator
        numerator = self.lambdaVal;
        if inputBigram in self.obsVoc:
            numerator = self.lambdaVal + self.obsVoc[inputBigram];
        
        # calculate denominator
        if inputBigram[0] == self.BEGIN_ARTICLE:
            sampleSize = self.numArticles; # frequency of virtual "begin-article" event
        else:
            sampleSize = self.uniObsVoc[inputBigram[0]]; # frequency of conditioning word (unigram) in test set
        denominator = sampleSize + self.lambdaVal*self.VOCAB_SIZE;
        
        # return discounted probability
        return numerator / denominator;
    
    '''
    Returns number of instances of given bigram in train set
    '''
    def getFreq(self, bigram):
        
        # check if bigram does not appear in train set
        if self.obsVoc.get(bigram) is None:
            return 0;
        
        # return bigram frequency
        return self.obsVoc[bigram];
    
    '''
    Returns number of instances of a bigram where conditioning word
    is the virtual "begin-article" event.
    This bigram appears only at the beginning, therefore returns 1
    '''
    def getVirtualFreq(self):
        return 1;
    
    '''
    Returns discounted probability of a bigram where conditioning word
    is the virtual "begin-article" event.
    The only bigram with "begin-article" as the conditioning word is at
    the beginning of the document, therefore sample size is 1
    '''
    def getVirtualDiscProb(self):
        
        # calculate numerator
        numerator = self.getVirtualFreq() + self.lambdaVal;
        
        # calculate denominator
        denominator = 1 + self.lambdaVal*self.VOCAB_SIZE;
        
        # return discounted probability
        return numerator / denominator;
    